package gm

import (
	"github.com/emmansun/gmsm/sm4"
)

func EncryptSm4Ecb(key []byte, data []byte) []byte {
	cipher, err := sm4.NewCipher(key)
	if err != nil {
		return nil
	}

	c := (len(data) + sm4.BlockSize) / sm4.BlockSize
	size := c * sm4.BlockSize
	pad := size - len(data)

	encrypted := make([]byte, size)

	i, j := 0, sm4.BlockSize
	for ; j <= len(data); i, j = j, j+sm4.BlockSize {
		cipher.Encrypt(encrypted[i:j], data[i:j])
	}

	if pad > 0 {
		last := make([]byte, 0, sm4.BlockSize)

		if i < len(data) {
			last = append(last, data[i:]...)
		}

		for p := len(data) - i; p < sm4.BlockSize; p = p + 1 {
			last = append(last, byte(pad))
		}

		cipher.Encrypt(encrypted[i:], last)
	}

	return encrypted
}

func DecryptSm4Ecb(key []byte, encrypted []byte) []byte {
	cipher, err := sm4.NewCipher(key)
	if err != nil {
		return nil
	}

	data := make([]byte, len(encrypted))

	i, j := 0, sm4.BlockSize
	for ; i < len(encrypted); i, j = j, j+sm4.BlockSize {
		cipher.Decrypt(data[i:j], encrypted[i:j])
	}
	pad := data[len(data)-1]
	return data[0 : i-int(pad)]
}
